import os
from naive_bayes import NaiveBayesClassifier

def load_dataset(data_dir):
    """加载数据集（读取类别文件夹下的文本）"""
    texts, labels = [], []
    # 遍历每个类别文件夹
    for label in os.listdir(data_dir):
        label_path = os.path.join(data_dir, label)
        if not os.path.isdir(label_path):
            continue  # 跳过非文件夹
        # 读取文件夹内所有文件
        for filename in os.listdir(label_path):
            file_path = os.path.join(label_path, filename)
            try:
                with open(file_path, 'r', encoding='latin-1') as f:
                    texts.append(f.read())
                    labels.append(label)
            except:
                continue  # 忽略读取失败的文件
    return texts, labels

def main():
    # 获取当前文件所在目录的绝对路径（解决相对路径问题）
    current_dir = os.path.dirname(os.path.abspath(__file__))
    # 拼接数据集路径（确保data文件夹与run_bayes.py同目录）
    train_dir = os.path.join(current_dir, "data", "20news-bydate-train")
    test_dir = os.path.join(current_dir, "data", "20news-bydate-test")

    # 加载训练集和测试集（仅保留2个类别，加速测试）
    target_classes = ['alt.atheism', 'comp.graphics']
    train_texts, train_labels = load_dataset(train_dir)
    test_texts, test_labels = load_dataset(test_dir)

    # 过滤出目标类别
    train_texts = [t for t, l in zip(train_texts, train_labels) if l in target_classes]
    train_labels = [l for l in train_labels if l in target_classes]
    test_texts = [t for t, l in zip(test_texts, test_labels) if l in target_classes]
    test_labels = [l for l in test_labels if l in target_classes]

    # 打印数据集信息
    print(f"训练集样本数：{len(train_texts)}（类别：{target_classes}）")
    print(f"测试集样本数：{len(test_texts)}")

    # 训练模型
    clf = NaiveBayesClassifier()
    clf.train(train_texts, train_labels)

    # 评估模型
    accuracy = clf.evaluate(test_texts, test_labels)
    print(f"模型准确率：{accuracy:.4f}")

    # 示例预测
    sample = "The graphics card renders 3D models efficiently."
    print(f"\n预测样本：{sample}")
    print(f"预测类别：{clf.predict(sample)}")

if __name__ == "__main__":
    main()